from mindquantum import *
from tqdm import trange
import numpy as np
import matplotlib.pyplot as plt

def s2(density_matrix):
    rho2 = np.matmul(density_matrix, density_matrix)
    return -np.log2(np.trace(rho2)).real

def get_rs(rho, list_qubits2keep):
    list_qubits2traceout = list(set(range(int(np.log2(len(rho)))))-set(list_qubits2keep))
    counter_qubits = int(np.log2(len(rho)))
    rho_tensor = rho.reshape((2 ,2)*counter_qubits)
    for i in list_qubits2traceout:
        index2trace = i - list_qubits2traceout.index(i)
        rho_tensor = np.trace(rho_tensor, axis1 = counter_qubits -1 - index2trace, axis2 = 2*counter_qubits -1 - index2trace)
        counter_qubits += -1
    rho_rs = rho_tensor.reshape(2**counter_qubits, 2**counter_qubits)
    return rho_rs
def s2_ising(rho, n = 8):
    res = 0
    for ii in range(n-1):
        res += s2(get_rs(rho, list(set(list(range(n))) - set([ii,ii+1]))))
    return res
def grad_s2_ising(circ, params, n=8, step_length = 1e-7):
    sim = Simulator('projectq', circ.n_qubits)
    grad = np.zeros(len(params))
    for jj in range(len(params)):
        delta = np.zeros(len(params))
        delta[jj] = step_length
        sim.reset()
        sim.apply_circuit(circ, params + delta)
        qs = sim.get_qs()
        rho = np.outer(qs, qs.conjugate())
        s2_1 = s2_ising(rho)
        sim.reset()
        sim.apply_circuit(circ, params)
        qs = sim.get_qs()
        rho = np.outer(qs, qs.conjugate())
        s2_0 = s2_ising(rho)
        grad[jj] = (s2_1-s2_0)/(1e-5)
    return grad
def noise_generator(params, sigma):
    params = np.random.randn(len(params))*sigma + params
    return params

#定义随机Pauli旋转门
def RP(rotation_angle):
    a = np.random.randint(0, 3)
    if a == 0:
        return RX(rotation_angle)
    elif a == 1:
        return RY(rotation_angle)
    elif a == 2:
        return RZ(rotation_angle)
    else: 
        print("error in random Pauli gates")
#定义barren plateau的ansatz，线路图参考readme或者main.ipynb文件
def bpansatz(n,p):
    par = np.array(range(n*p)).astype(np.str)
    init_state_circ = UN(RY(np.pi/4),n)
    qc = Circuit() + init_state_circ
    for j in range(p):
        for i in range(n):
            theta = par[i*p+j]
            qc += RP(theta).on(i)
        for ii in range(n-1):
            qc += Z(ii,ii+1)
    return qc

if __name__ == '__main__':
    n = 8
    bpcircuit = bpansatz(n, 100)
    op_ising = QubitOperator()
    for i in range(n-1):
        op_ising += QubitOperator(f'Z{i} Z{i+1}')
    ham_ising = Hamiltonian(op_ising)
    params_set = np.random.rand(len(bpcircuit.params_name))*0.1

    eta = 0.01                                                          #learning rate
    sigma = 0.015                                                       #error variance
    alpha = 1                                                           #learning rate of S2
    n_steps = 200                                                       #number of training steps
    params = params_set
    sim = Simulator('projectq', bpcircuit.n_qubits)
    grad_ops = sim.get_expectation_with_grad(ham_ising, bpcircuit)
    cost = np.zeros(n_steps)
    s2_list = np.zeros(n_steps)
    for i in trange(n_steps):
        sim.reset()
        sim.apply_circuit(bpcircuit, params)
        qs = sim.get_qs()
        rho = np.outer(qs, qs.conjugate())
        s2_list[i] = s2_ising(rho,n)
        cost[i] = (grad_ops(params)[0][0,0]).real
        if s2_list[i]>4:
            params = params - eta*((grad_ops(params)[1][0][0]) + alpha*grad_s2_ising(bpcircuit, params, n))
        else:
            params = params - eta*(grad_ops(params)[1][0][0])
        params = noise_generator(params, sigma)
    print('minimum with S2 descent:', cost.min())
    plt.plot(cost, label = 'noisy circuit with descent S2')
    plt.plot(s2_list, label = 'noisy S2 with descent S2')
    plt.legend()

    params = params_set
    cost = np.zeros(n_steps)
    s2_list = np.zeros(n_steps)

    for i in trange(n_steps):
        sim.reset()
        sim.apply_circuit(bpcircuit, params)
        qs = sim.get_qs()
        rho = np.outer(qs, qs.conjugate())
        s2_list[i] = s2_ising(rho, n)
        cost[i] = (grad_ops(params)[0][0,0]).real
        params = params - eta*(grad_ops(params)[1][0][0])
        params = noise_generator(params, sigma)
    print('minimum of the conventional method', cost.min())
    plt.plot(cost, label = 'noisy circuit')
    plt.plot(s2_list, label = 'noisy S2')
    plt.legend()
    plt.title('Learning curve of Ising Hamiltonian')
    plt.savefig('train_nisq.svg')